__all__=["img_to_pts","fit_ellipse","least_square_circle","solve_drogue_wa","solve_drogue"]

import cv2
import numpy as np
from scipy.optimize import least_squares
import itertools


def img_to_pts(img_bin, min_area = 15, max_area = 1000):
    '''提取图片中的特征点'''
    # 连通域分析
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img_bin, connectivity=4)

    gt_pts = []
    area_list = []
    for stat,centroid in zip(stats[1:],centroids[1:]):
 
        if stat[-1] < min_area or stat[-1] > max_area:
            continue
        gt_pts.append(centroid)
        area_list.append(stat[-1])
    return gt_pts,area_list


def fit_ellipse(x,y):
    l = np.size(x,axis=0)
    D1 = np.vstack([(x*x),(x*y),(y*y)]).T
    D2 = np.vstack([x,y,np.ones((1,np.size(x,axis=1)))]).T
    S1 = np.matmul(D1.T,D1)
    S2 = np.matmul(D1.T,D2)
    S3 = np.matmul(D2.T,D2)
    T = -1*np.matmul(np.linalg.pinv(S3),S2.T)
    M = S1+np.matmul(S2,T)
    M = np.vstack([M[2,:]/2,-1*M[1,:],M[0,:]/2])
    egienvalue, featurevec = np.linalg.eig(M)
    cond = 4*featurevec[0,:]*featurevec[2,:]-featurevec[1,:]*featurevec[1,:]
    vaild_vec1 = featurevec[:,cond>0]
    vaild_vec = np.vstack([vaild_vec1,np.matmul(T,vaild_vec1)])
    
    if vaild_vec.size == 0 or vaild_vec[0].size > 1:
        return 0,0,0,0,0,False
    else:
        A = vaild_vec[0]
        B = vaild_vec[1]
        C = vaild_vec[2]
        D = vaild_vec[3]
        E = vaild_vec[4]
        F = vaild_vec[5]
        if A<0:
            A=A*(-1)
            B=B*(-1)
            C=C*(-1)
            D=D*(-1)
            E=E*(-1)
            F=F*(-1)

        h = (B*E-2*C*D)/(4*A*C-B*B)  
        k = (B*D-2*A*E)/(4*A*C-B*B)
        a = np.sqrt(2*((A*E*E-B*D*E+C*D*D)/(4*A*C-B*B)-F)/(A+C-np.sqrt((A-C)*(A-C)+B*B)))
        b = np.sqrt(2*((A*E*E-B*D*E+C*D*D)/(4*A*C-B*B)-F)/(A+C+np.sqrt((A-C)*(A-C)+B*B))) 
  
        theta = 0.5*np.arctan2(B,A-C)

        return h,k,a,b,theta,True


def least_square_circle(pts):
    X = pts[:,0]
    Y = pts[:,1]
    A = np.vstack([X,Y,np.ones_like(X)]).T
    B = X*X + Y*Y
    m_vec = np.linalg.pinv(A).dot(B)/2
    cx,cy,m3 = m_vec
    r = np.sqrt(cx*cx + cy*cy + 2*m3)
    return cx,cy,r

# def solve_drogue(taper_px_ptL,cam_K,taper_radius):
#     result_dict = {
#         "rpe": np.nan,
#         "avaiabel_pt_num": 0,
#         "t_vec": np.array([np.nan]*3),
#         "img_valid": False,
#         "drogue_px_calcL": np.array([]),
#     }
#     N_total = len(taper_px_ptL)
#     N = 6 #选取6个特征点用于解算
#     pt_num = len(taper_px_ptL)
#     # 特征点不足以解算
#     if pt_num < N:
#        
#         return result_dict
#     taper_px_pts = np.array(taper_px_ptL)
#     data_dictL = []
#     # 拟合圆并按照在圆上的最大点数和重投影误差排序
#     for pts in itertools.combinations(taper_px_ptL,N):
#         cx,cy,r = least_square_circle(np.array(pts))
#         circle_center = np.array([cx,cy])
#         rpes = np.sqrt(np.abs(np.sum((taper_px_pts-circle_center)*(taper_px_pts-circle_center),axis=1) - r*r)/N_total/r)
#         ind = rpes<0.7
#         rpes_valid = rpes[ind]
#         available_pt_num = np.size(rpes_valid)   #看这个圆的拟合程度
#         if available_pt_num == 0:
#             rpe = np.nan
#         else:
#             rpe = np.average(rpes_valid)

#         # 均为降序则取负
#         data_dict = {
#             "available_pt_num": available_pt_num,
#             "rpe": -rpe,
#             "cx":cx,
#             "cy":cy,
#             "r":r,
#             "drogue_px_calcL":pts
#         }

#         data_dictL.append(data_dict)
        
#     # 取重投影误差最小的
#     data_dictL_sorted = sorted(data_dictL,key=lambda x: [x["available_pt_num"],x["rpe"]],reverse=True)
#     data_dict = data_dictL_sorted[0]
# 
#     # 有效点数目足够
#     if data_dict["available_pt_num"] >= N:
#         depth = taper_radius*cam_K[0][0]/data_dict["r"]
#         t_vec = depth*np.linalg.inv(cam_K).dot(np.array([data_dict["cx"],data_dict["cy"],1]))
#         result_dict["t_vec"] = t_vec
#         result_dict["rpe"] = -1 * data_dict["rpe"]
#         result_dict["available_pt_num"] = data_dict["available_pt_num"]
#         result_dict["img_valid"] = True
#         result_dict["drogue_px_calcL"] = data_dict["drogue_px_calcL"]
#     return result_dict

#鱼眼镜头
def solve_drogue_wa(taper_px_ptL,mapping_coeffs, distort_center, stretch_matrix,taper_radius):
    result_dict = {
        "rpe": np.nan,
        "avaiabel_pt_num": 0,
        "t_vec": np.array([np.nan]*3),
        "img_valid": False,
        "drogue_px_calcL": np.array([]),
    }
    N = 6 #选取6个特征点用于解算
    pt_num = len(taper_px_ptL)
    # 特征点不足以解算
    if pt_num < N:
        return result_dict

    data_dictL = []
    # 拟合圆并按照在圆上的最大点数和重投影误差排序
    for pts in itertools.combinations(taper_px_ptL,N):
        pts_T = np.array(pts).T
        pts_undistort = (np.linalg.inv(stretch_matrix)).dot(pts_T-distort_center.T)
        rho = np.sqrt(pts_undistort[0,:]**2 + pts_undistort[1,:]**2)
        pts_3D = np.vstack((pts_undistort[0,:],pts_undistort[1,:],np.sum((mapping_coeffs.T)*(np.vstack((np.array([1]*len(rho)),rho**2,rho**3,rho**4))),axis=0)))/1000
        pts_3D_x = pts_3D[0,:].T  #单位为m
        pts_3D_y = pts_3D[1,:].T
        pts_3D_x = pts_3D_x[np.newaxis, :]
        pts_3D_y = pts_3D_y[np.newaxis, :]
        
        cx,cy,a,b,theta,flag = fit_ellipse(np.array(pts_3D_x),np.array(pts_3D_y))

        if flag:
            rpe = np.max(np.sqrt(np.abs((((pts_3D_x-cx)*np.cos(theta))*((pts_3D_x-cx)*np.cos(theta))+(((pts_3D_y-cy)*np.sin(theta))*((pts_3D_y-cy)*np.sin(theta))))/a/a + \
                    (((pts_3D_x-cx)*np.sin(theta))*((pts_3D_x-cx)*np.sin(theta))+(((pts_3D_y-cy)*np.cos(theta))*((pts_3D_y-cy)*np.cos(theta))))/b/b - 1))/a)
            

            # 均为降序则取负
            data_dict = {
                "rpe": rpe,
                "cx":cx,
                "cy":cy,
                "a":a,
            }
            data_dictL.append(data_dict)
        
    if np.size(data_dictL) > 0:
        # 取重投影误差最小的
        data_dictL_sorted = sorted(data_dictL,key=lambda x: x["rpe"],reverse=False)
        data_dict = data_dictL_sorted[0]
        # print("drogue_all", data_dictL_sorted)
        # print("drogue",data_dict["rpe"])
        # 有效点数目足够
        if data_dict["rpe"] < 10:
            a = data_dict["a"]
            cx = data_dict["cx"]
            cy = data_dict["cy"]   #they are all lists whose length is 1
            lamb = taper_radius/a[0]
            t_vec = lamb * np.array([cx[0],cy[0],1])
            result_dict["t_vec"] = t_vec
            result_dict["rpe"] =  data_dict["rpe"]
            result_dict["img_valid"] = True
        return result_dict
    else:
        return result_dict


# # ellipse fit
# def solve_drogue(taper_px_ptL,cam_K,taper_radius):
#     result_dict = {
#         "rpe": np.nan,
#         "avaiabel_pt_num": 0,
#         "t_vec": np.array([np.nan]*3),
#         "img_valid": False,
#         "drogue_px_calcL": np.array([]),
#     }
#     N = 6  #选取6个特征点用于解算
#     pt_num = len(taper_px_ptL)
#     # 特征点不足以解算
#     if pt_num < N:
#      
#         return result_dict
#     taper_px_pts = np.array(taper_px_ptL)
#     data_dictL = []
#     # 拟合圆并按照在圆上的最大点数和重投影误差排序
#     for pts in itertools.combinations(taper_px_ptL,N):
#         pts_array = np.array(pts)
#         x = pts_array[:,0]
#         y = pts_array[:,1]
#         a, b, h, k, theta = fit_ellipse(x,y)
#         rpes = np.sqrt(np.abs(((x-h)*np.cos(theta)/a)**2+(((y-k)*np.sin(theta)/b)**2)-1))
#         ind = rpes<100
#         rpe = np.mean(rpes)

#         # 均为降序则取负
#         data_dict = {
#             "rpe": rpe,
#             "a":a,
#             "b":b,
#             "h":h,
#             "k":k,
#             "theta":theta,
#             "drogue_px_calcL":pts
#         }

#         data_dictL.append(data_dict)
        
#     # 取重投影误差最小的
#     data_dictL_sorted = sorted(data_dictL,key=lambda x: x["rpe"],reverse=False)
#     data_dict = data_dictL_sorted[0]

#     # 有效点数目足够
#   
#     depth = taper_radius*cam_K[0][0]/data_dict["a"]
#     t_vec = depth*np.linalg.inv(cam_K).dot(np.array([data_dict["h"],data_dict["k"],1]))
#     result_dict["t_vec"] = t_vec
#     result_dict["rpe"] = -1 * data_dict["rpe"]
#     result_dict["img_valid"] = True
#     result_dict["drogue_px_calcL"] = data_dict["drogue_px_calcL"]
#     return result_dict

def solve_drogue(taper_px_ptL,cam_K,taper_radius):
    result_dict = {
        "rpe": np.nan,
        "avaiabel_pt_num": 0,
        "t_vec": np.array([np.nan]*3),
        "img_valid": False,
        "drogue_px_calcL": np.array([]),
    }
    N_total = len(taper_px_ptL)
    N = 6#选取6个特征点用于解算
    pt_num = len(taper_px_ptL)
    # 特征点不足以解算
    if pt_num < N:
        return result_dict

    data_dictL = []
    # 拟合圆并按照在圆上的最大点数和重投影误差排序
    for pts in itertools.combinations(taper_px_ptL,N):
 
        x = np.array(pts)[:,0]
        y = np.array(pts)[:,1]
        x = x[np.newaxis, :]
        y = y[np.newaxis, :]
        h,k,a,b,theta,flag = fit_ellipse(x,y)
        if flag:
            rpes = np.sqrt(np.abs((((x-h)*np.cos(theta))*((x-h)*np.cos(theta))+(((y-k)*np.sin(theta))*((y-k)*np.sin(theta))))/a/a + \
                        (((x-h)*np.sin(theta))*((x-h)*np.sin(theta))+(((y-k)*np.cos(theta))*((y-k)*np.cos(theta))))/b/b - 1))
            ind = rpes<0.3
            rpes_valid = rpes[ind]
            available_pt_num = np.size(rpes_valid)   #看这个圆的拟合程度
            if available_pt_num == 0:
                rpe = np.nan
            else:
                rpe = np.max(rpes_valid)
            if rpe > 0.3:
                available_pt_num = 0
            # 均为降序则取负
            data_dict = {
                "valid_num": available_pt_num,
                "rpe": -1*rpe,
                "cx":h,
                "cy":k,
                "a":max(a,b),
                "drogue_px_calcL":pts
            }
            data_dictL.append(data_dict)
    if np.size(data_dictL) > 0:
        # 取重投影误差最小的
        data_dictL_sorted = sorted(data_dictL,key=lambda x: [x["valid_num"], x["rpe"]],reverse=True)
        data_dict = data_dictL_sorted[0]

        # print("drogue",data_dict["rpe"],data_dict["valid_num"])
        # if data_dict["valid_num"] >= N:
        if data_dict["valid_num"] >= 4:
            
            depth = taper_radius*cam_K[0][0]/data_dict["a"]
            t_vec = depth*np.linalg.inv(cam_K).dot(np.array([data_dict["cx"][0],data_dict["cy"][0],1]))
            result_dict["t_vec"] = t_vec
            result_dict["rpe"] = -1*data_dict["rpe"]
            result_dict["img_valid"] = True
            result_dict["drogue_px_calcL"] = data_dict["drogue_px_calcL"]
            # print(t_vec)
        return result_dict
    else:
        return result_dict

